
import cv2
import helper
import numpy as np
import rasterio


# 一般地，3 bands 视为 R G B，4 bands 视为 R G B NIR
def main(args):
    dst = args.src + '.png'

    with rasterio.open(args.src) as dataset:
        if dataset.count == 1:
            arr = dataset.read(dataset.indexes[0])
            cv2.imwrite(dst, arr)

        if dataset.count >= 3:
            if args.use_nir:
                if args.bgr_index is not None:
                    ics = map(int, args.bgr_index.split(','))
                else:
                    ics = (3, 2, 1, 4)
            else:
                if args.bgr_index is not None:
                    ics = map(int, args.bgr_index.split(','))
                else:
                    ics = (3, 2, 1)

            # (channels, height, width) -> (height, width, channels)
            arr = dataset.read(list(ics)).transpose((1, 2, 0))
            # print(arr.shape)

            if not args.use_nir and arr.shape[2] == 4:
                # (height, width, channels=4) -> (height, width, channels=3)
                arr = cv2.cvtColor(arr, cv2.COLOR_BGRA2BGR)
                # print(arr.shape)

            cv2.imwrite(dst, arr)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src')
    parser.add_argument('--bgr-index')
    parser.add_argument('--use-nir', action='store_true')
    args = parser.parse_args()
    main(args)
